from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.engine.url import URL
from libs.dbutils import sqlalchemy_config as settings


class DBHelper():
    '''
        @说明    封装的sqlalchemy操作相关
        @时间    2020-04-01 15:34:49
        @作者    gzt
    '''

    # 数据模型继承的基类
    BASE = declarative_base()

    def __init__(self, user_name=None, passwrod=None, ip_address=None, port=None, db_name=None, *args, **kwargs):
        self.user_name = user_name
        if self.user_name is None:
            self.user_name = settings.USER_NAME

        self.passwrod = passwrod
        if self.passwrod is None:
            self.passwrod = settings.PASSWORD

        self.ip_address = ip_address

        if self.ip_address is None:
            self.ip_address = settings.IP_ADDRESS

        self.port = port

        if self.port is None:
            self.port = settings.PORT

        self.db_name = db_name

        if self.db_name is None:
            self.db_name = settings.DATABASE

        # 创建连接对象
        self.engine = create_engine(URL(**{'database': self.db_name,
                                           'username': self.user_name,
                                           'password': self.passwrod,
                                           'host': self.ip_address,
                                           'port': self.port,
                                           'drivername': 'mysql+pymysql'}), encoding='utf-8', convert_unicode=True)

        self.session = None

    def __enter__(self):
        '''
            说  明     :初始化连接
            时  间     :2020-04-01 15:53:18
            作  者     :gzt
        '''
        Session = sessionmaker()
        Session.configure(bind=self.engine)

        self.session = Session()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        '''
            说  明     :退出后关闭数据库的连接
            时  间     :2020-04-01 15:58:55
            作  者     :gzt
        '''

        self.session.close()

    @classmethod
    def create_all(cls):
        '''
            说  明     :创建所有的数据库表
            时  间     :2020-04-01 15:59:46
            作  者     :gzt
        '''
        cls.BASE.metadata.create_all(bind=cls().engine)

    def add(self, entity):
        '''
            说  明     :添加实体
            返  回     : None: 添加失败
                        entity:添加成功后的实体
            时  间     :2020-04-01 16:01:53
            作  者     :gzt
        '''
        try:
            self.session.add(entity)
            self.session.commit()
            return entity
        except SQLAlchemyError as e:
            return None
